﻿// 2024/12/18: 首个版本
// 2025/4/8: 修正yaw angle解析
// 2025/8/4: 不使用cat宏
// 2025/8/4: 改为继承SessionSample。改回单头文件

#ifndef POINT_CLOUD_V2_H
#define POINT_CLOUD_V2_H

#include "spadas.h"

namespace point_cloud_v2
{
	using namespace spadas;

	struct PointCloudData
	{
		UInt pointCount; // 点个数
		FloatMat points; // Coordinates of points, the size of array is Nx3, while N is number of points, 3 means x(front),y(left),z(up), in meters / 各点坐标，数组尺寸为Nx3，N为个数，3分别为x(前),y(左),z(上)轴坐标，单位为米
		Array<Float> intensities; // Intensity of points, in percentages, can be empty / 各点强度，单位为%，可为空
		Array<ColorRGB> colors; // Color of points, can be empty / 各点颜色，可为空
		Array<Word> laserIDs; // Laser ID of points, can be empty (For example, from bottom to top, ID of 32 line lidar are 0~31) / 各点的激光束ID，可为空（如32线雷达从下往上ID分别为0~31）
		Array<Float> timeOffsets; // Time offset (relative to sample's timestamp) of points, in microseconds, can be empty / 各点相对于样本时间偏置，单位微秒，可为空
		Array<Float> yawAngles; // Yaw angle of points, in degrees, can be empty / 各点横摆角，单位度，可为空
		Array<Byte> mirrorIDs; // Mirror ID of points, can be empty / 各点镜面ID，可为空
		Map<IntKey, Array<Byte> > byteValues; // Dictionary for byte values. The key is field ID, the value is byte values / Byte值字典，键为字段ID，值为Byte值
		Map<IntKey, Array<Short> > shortValues; // Dictionary for short values. The key is field ID, the value is short values / Short值字典，键为字段ID，值为Short值
		Map<IntKey, Array<Float> > floatValues; // Dictionary for float values. The key is field ID, the value is float values / Float值字典，键为字段ID，值为Float值

		PointCloudData()
		{
			pointCount = 0;
		}
	};

	class PointCloud : public SessionSample
	{
	private:
		class PointCloudVars : public SessionSampleVars
		{
		public:
			SPADAS_VARS(PointCloud, SessionSampleVars)
			PointCloudData data;

			SessionMatrixSample toMatrixSample()
			{
				UInt iCol = UINF, rCol = UINF, gCol = UINF, bCol = UINF, lidCol = UINF, toCol = UINF, yaCol = UINF, midCol = UINF;
				Array<UInt> byteFieldColTable(100, UINF); // 100~199
				Array<UInt> shortFieldColTable(100, UINF); // 200~299
				Array<UInt> floatFieldColTable(100, UINF); // 300~399

				UInt nPoints = data.pointCount;
				if (data.points.nDims() != 2 || data.points.dimAt(1) != 3) nPoints = 0;
				else if (nPoints > data.points.dimAt(0)) nPoints = data.points.dimAt(0);

				UInt nCols = 3;
				if (nPoints > 0)
				{
					if (data.intensities.size() == nPoints) iCol = nCols++;
					if (data.colors.size() == nPoints)
					{
						rCol = nCols++;
						gCol = nCols++;
						bCol = nCols++;
					}
					if (data.laserIDs.size() == nPoints) lidCol = nCols++;
					if (data.timeOffsets.size() == nPoints) toCol = nCols++;
					if (data.yawAngles.size() == nPoints) yaCol = nCols++;
					if (data.mirrorIDs.size() == nPoints) midCol = nCols++;
					for (auto pair = data.byteValues.keyValues().firstElem(); pair.valid(); ++pair)
					{
						Int fieldID = pair->key.value();
						if (fieldID < 100 || fieldID >= 200 || pair->value.size() != nPoints) continue;
						byteFieldColTable[fieldID - 100] = nCols++;
					}
					for (auto pair = data.shortValues.keyValues().firstElem(); pair.valid(); ++pair)
					{
						Int fieldID = pair->key.value();
						if (fieldID < 200 || fieldID >= 300 || pair->value.size() != nPoints) continue;
						shortFieldColTable[fieldID - 200] = nCols++;
					}
					for (auto pair = data.floatValues.keyValues().firstElem(); pair.valid(); ++pair)
					{
						Int fieldID = pair->key.value();
						if (fieldID < 300 || fieldID >= 400 || pair->value.size() != nPoints) continue;
						floatFieldColTable[fieldID - 300] = nCols++;
					}
				}

				SessionMatrixSample sample(Size2D::dims(nPoints + 1, nCols));
				sample.timestamp() = timestamp;

				Float *headData = sample.matData().data();
				Float *pointData = &sample.matData()[nCols];
				
				Float *xyzData = data.points.data();
				UInt k = 0;
				for (UInt a = 0; a < nPoints; a++)
				{
					pointData[a * nCols] = xyzData[k++];
					pointData[a * nCols + 1] = xyzData[k++];
					pointData[a * nCols + 2] = xyzData[k++];
				}

				if (iCol != UINF)
				{
					headData[iCol] = 1;
					Float *iData = data.intensities.data();
					for (UInt a = 0; a < nPoints; a++)
					{
						pointData[a * nCols + iCol] = iData[a];
					}
				}

				if (rCol != UINF)
				{
					headData[rCol] = 2;
					headData[gCol] = 3;
					headData[bCol] = 4;
					ColorRGB *colorData = data.colors.data();
					for (UInt a = 0; a < nPoints; a++)
					{
						pointData[a * nCols + rCol] = colorData[a].r;
						pointData[a * nCols + gCol] = colorData[a].g;
						pointData[a * nCols + bCol] = colorData[a].b;
					}
				}

				if (lidCol != UINF)
				{
					headData[lidCol] = 5;
					Word *lidData = data.laserIDs.data();
					for (UInt a = 0; a < nPoints; a++)
					{
						pointData[a * nCols + lidCol] = lidData[a];
					}
				}

				if (toCol != UINF)
				{
					headData[toCol] = 6;
					Float *toData = data.timeOffsets.data();
					for (UInt a = 0; a < nPoints; a++)
					{
						pointData[a * nCols + toCol] = toData[a];
					}
				}

				if (yaCol != UINF)
				{
					headData[yaCol] = 7;
					Float *yaData = data.yawAngles.data();
					for (UInt a = 0; a < nPoints; a++)
					{
						pointData[a * nCols + yaCol] = yaData[a];
					}
				}

				if (midCol != UINF)
				{
					headData[midCol] = 8;
					Byte *midData = data.mirrorIDs.data();
					for (UInt a = 0; a < nPoints; a++)
					{
						pointData[a * nCols + midCol] = midData[a];
					}
				}

				for (UInt i = 0; i < byteFieldColTable.size(); i++)
				{
					if (byteFieldColTable[i] == UINF) continue;
					UInt byteFieldCol = byteFieldColTable[i];
					headData[byteFieldCol] = 100 + i;
					Byte *byteFieldData = data.byteValues[IntKey(100 + i)].data();
					for (UInt a = 0; a < nPoints; a++)
					{
						pointData[a * nCols + byteFieldCol] = byteFieldData[a];
					}
				}

				for (UInt i = 0; i < shortFieldColTable.size(); i++)
				{
					if (shortFieldColTable[i] == UINF) continue;
					UInt shortFieldCol = shortFieldColTable[i];
					headData[shortFieldCol] = 200 + i;
					Short *shortFieldData = data.shortValues[IntKey(200 + i)].data();
					for (UInt a = 0; a < nPoints; a++)
					{
						pointData[a * nCols + shortFieldCol] = shortFieldData[a];
					}
				}

				for (UInt i = 0; i < floatFieldColTable.size(); i++)
				{
					if (floatFieldColTable[i] == UINF) continue;
					UInt floatFieldCol = floatFieldColTable[i];
					headData[floatFieldCol] = 300 + i;
					Float *floatFieldData = data.floatValues[IntKey(300 + i)].data();
					for (UInt a = 0; a < nPoints; a++)
					{
						pointData[a * nCols + floatFieldCol] = floatFieldData[a];
					}
				}

				return sample;
			}

			Bool fromSample(String protocol, SessionSample sample) override
			{
				SessionMatrixSample matrixSample;
				if (!sample.is<SessionMatrixSample>(matrixSample)) return FALSE;

				timestamp = matrixSample.timestamp();
				
				const UInt nRows = matrixSample.rows();
				const UInt nCols = matrixSample.cols();
				if (nCols < 3) return FALSE;

				Float *pointData = NULL;
				UInt nPoints = 0;
				UInt iCol = UINF, rCol = UINF, gCol = UINF, bCol = UINF, lidCol = UINF, toCol = UINF, yaCol = UINF, midCol = UINF;
				Bool hasColor = FALSE;
				Array<UInt> byteFieldColTable(100, UINF); // 100~199
				Array<UInt> shortFieldColTable(100, UINF); // 200~299
				Array<UInt> floatFieldColTable(100, UINF); // 300~399

				if (protocol == "point-cloud-v2")
				{
					if (nRows == 0) return FALSE;

					pointData = matrixSample.matData().data();
					nPoints = nRows - 1;

					for (UInt a = 3; a < nCols; a++)
					{
						UInt fieldType = (UInt)pointData[a];
						switch (fieldType)
						{
						case 1:
							iCol = a;
							break;
						case 2:
							rCol = a;
							break;
						case 3:
							gCol = a;
							break;
						case 4:
							bCol = a;
							break;
						case 5:
							lidCol = a;
							break;
						case 6:
							toCol = a;
							break;
						case 7:
							yaCol = a;
							break;
						case 8:
							midCol = a;
							break;
						default:
							if (fieldType >= 100 && fieldType < 200)
							{
								byteFieldColTable[fieldType - 100] = a;
							}
							else if (fieldType >= 200 && fieldType < 300)
							{
								shortFieldColTable[fieldType - 200] = a;
							}
							else if (fieldType >= 300 && fieldType < 400)
							{
								floatFieldColTable[fieldType - 300] = a;
							}
							break;
						}
					}

					pointData = &pointData[nCols];
					hasColor = rCol != UINF && gCol != UINF && bCol != UINF;
				}
				else if (protocol == "point-cloud-v1")
				{
					pointData = matrixSample.matData().data();
					nPoints = nRows;
				}
				else return FALSE;

				data.pointCount = nPoints;
				data.points = nPoints > 0 ? FloatMat(nPoints, 3) : FloatMat();
				data.intensities = Array<Float>();
				data.colors = Array<ColorRGB>();
				data.laserIDs = Array<Word>();
				data.timeOffsets = Array<Float>();
				data.yawAngles = Array<Float>();
				data.mirrorIDs = Array<Byte>();
				data.byteValues = Map<IntKey, Array<Byte> >();
				data.shortValues = Map<IntKey, Array<Short> >();
				data.floatValues = Map<IntKey, Array<Float> >();

				if (nPoints > 0)
				{
					Float *xyzData = data.points.data();
					UInt k = 0;
					for (UInt a = 0; a < nPoints; a++)
					{
						xyzData[k++] = pointData[a * nCols];
						xyzData[k++] = pointData[a * nCols + 1];
						xyzData[k++] = pointData[a * nCols + 2];
					}

					if (iCol != UINF)
					{
						data.intensities = Array<Float>(nPoints);
						Float *iData = data.intensities.data();
						for (UInt a = 0; a < nPoints; a++)
						{
							iData[a] = pointData[a * nCols + iCol];
						}
					}

					if (hasColor)
					{
						data.colors = Array<ColorRGB>(nPoints);
						ColorRGB *colorData = data.colors.data();
						for (UInt a = 0; a < nPoints; a++)
						{
							colorData[a].r = (Byte)pointData[a * nCols + rCol];
							colorData[a].g = (Byte)pointData[a * nCols + gCol];
							colorData[a].b = (Byte)pointData[a * nCols + bCol];
						}
					}

					if (lidCol != UINF)
					{
						data.laserIDs = Array<Word>(nPoints);
						Word *lidData = data.laserIDs.data();
						for (UInt a = 0; a < nPoints; a++)
						{
							lidData[a] = (Word)pointData[a * nCols + lidCol];
						}
					}

					if (toCol != UINF)
					{
						data.timeOffsets = Array<Float>(nPoints);
						Float *toData = data.timeOffsets.data();
						for (UInt a = 0; a < nPoints; a++)
						{
							toData[a] = pointData[a * nCols + toCol];
						}
					}

					if (yaCol != UINF)
					{
						data.yawAngles = Array<Float>(nPoints);
						Float *yaData = data.yawAngles.data();
						for (UInt a = 0; a < nPoints; a++)
						{
							yaData[a] = pointData[a * nCols + yaCol];
						}
					}

					if (midCol != UINF)
					{
						data.mirrorIDs = Array<Byte>(nPoints);
						Byte *midData = data.mirrorIDs.data();
						for (UInt a = 0; a < nPoints; a++)
						{
							midData[a] = (Byte)pointData[a * nCols + midCol];
						}
					}

					for (UInt x = 0; x < 100; x++)
					{
						if (byteFieldColTable[x] != UINF)
						{
							UInt byteFieldCol = byteFieldColTable[x];
							Array<Byte> byteValues(nPoints);
							Byte *byteFieldData = byteValues.data();
							for (UInt a = 0; a < nPoints; a++)
							{
								byteFieldData[a] = (Byte)pointData[a * nCols + byteFieldCol];
							}
							data.byteValues[IntKey(100 + x)] = byteValues;
						}
					}

					for (UInt x = 0; x < 100; x++)
					{
						if (shortFieldColTable[x] != UINF)
						{
							UInt shortFieldCol = shortFieldColTable[x];
							Array<Short> shortValues(nPoints);
							Short *shortFieldData = shortValues.data();
							for (UInt a = 0; a < nPoints; a++)
							{
								shortFieldData[a] = (Short)pointData[a * nCols + shortFieldCol];
							}
							data.shortValues[IntKey(200 + x)] = shortValues;
						}
					}

					for (UInt x = 0; x < 100; x++)
					{
						if (floatFieldColTable[x] != UINF)
						{
							UInt floatFieldCol = floatFieldColTable[x];
							Array<Float> floatValues(nPoints);
							Float *floatFieldData = floatValues.data();
							for (UInt a = 0; a < nPoints; a++)
							{
								floatFieldData[a] = pointData[a * nCols + floatFieldCol];
							}
							data.floatValues[IntKey(300 + x)] = floatValues;
						}
					}
				}

				return TRUE;
			}
		};

	public:
		SPADAS_CLASS("point_cloud_v2.PointCloud", PointCloudVars)

		PointCloud() : SessionSample(new PointCloudVars)
		{}
		PointCloudData *operator ->()
		{
			return &var()->data;
		}

		static String getProtocolName(UInt channel)
		{
			return (String)"point-cloud-v2@" + channel;
		}
		SessionMatrixSample toMatrixSample()
		{
			return var()->toMatrixSample();
		}

		static Array<String> getProtocolNames(UInt channel)
		{
			return Array<String>::create(2, "point-cloud-v1", "point-cloud-v2").convert<String>([channel](auto& s){ return s + "@" + channel; });
		}
		Bool fromMatrixSample(String protocol, SessionMatrixSample matrixSample)
		{
			return fromSample(protocol, matrixSample);
		}
	};
}

#endif